﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;

using OptixDotNet;
using OptixDotNet.Math;
using OptixDotNet.Utils;

using RayDen.Library.Entity.Frames;
using RayDen.Optix.Engine.Components;
using RayDen.Optix.Engine.Core;
using RayDen.Optix.Engine.Managers;
using RayDen.Optix.Engine.Scene;
using RayDen.Optix.Engine.Scene.Entities;

namespace RayDen.Optix.Engine.Engines
{
    public class PathTracer : RenderEngine
    {

        private string lightPrefix = "lt";
        internal float yAvg = 0.1625f;
        internal float yMax = 2.5f;
        internal float apertureSize = 0.05f;

        private RdfManager sceneManager;

        internal PreethamSkyModel skyModel;
        private float SceneEpsilon = 1e-4f;
        private uint mFrame;



        public PathTracer()
        {
            this.CreateContext();
            sceneManager = new RdfManager(Context);
        }



        public override void OpenFrame(string framePath)
        {
            base.OpenFrame(framePath);
            sceneManager.OpenFrame(framePath);
            this.Scene = sceneManager.Scene;
        }

        public override void OpenFrame(FrameDescription frame)
        {
            sceneManager.OpenFrame(frame);
            this.Scene = sceneManager.Scene;
            this.Settings.ImageWidth = sceneManager.width;
            this.Settings.ImageHeight = sceneManager.height;
        }
        public override void Initialize(Action<BoundingBox> setCamera, params object[] parameters)
        {
            string rayGenPath = PtxDir + ("path_tracer.cu.ptx");
            string shaderPath = PtxDir + ("materials.cu.ptx");
            string skyShaderPath = PtxDir + ("sunsky.cu.ptx");
            string glassShaderPath = PtxDir + ("specular_mats.cu.ptx");

            var objFileName = parameters[0] as string;

            if (base.Context == null)
            {
                this.CreateContext();
            }

            var Context = base.Context;
            var Width = (int)parameters[1];
            var Height = (int)parameters[2];

            var desc = new BufferDesc()
            {
                Width = (uint)Width,
                Height = (uint)Height,
                Format = Format.Float3,
                Type = BufferType.InputOutput
            };

            var hdrBuffer = new OptixDotNet.Buffer(Context, desc);
            Context["hdr_buffer"].Set(hdrBuffer);

            Context["Y_log_av"].Set(yAvg);
            // maximum luminance
            Context["Y_max"].Set(yMax);

            Context["sig_e"].SetFloat3(new Vector3(0.7f, 0.7f, 0.7f));
            Context["sig_a"].SetFloat3(new Vector3(0.1f, 0.3f, 0.1f));

            /*-----------------------------------------------
             * Create the material that will be executed when there is an intersection
             *-----------------------------------------------*/

            Material material = new Material(Context);
            material.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "diffuse_all"); //glass_refract,diffuse
            material.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");


            /*-----------------------------------------------
             * Create the output buffer
             *-----------------------------------------------*/



            /*-----------------------------------------------
             * Create the ray-generation and exception programs
             *-----------------------------------------------*/
            Program rayGen = new Program(Context, rayGenPath, "pathtrace_camera");
            Program exception = new Program(Context, rayGenPath, "exception");
            Program miss = new Program(Context, shaderPath, "miss");
            Context["use_envmap"].Set(0u);


            //SETUP SKY
            if (skyModel == null)
            {
                skyModel = new PreethamSkyModel();
            }

            //Program skyMiss = new Program(Context, skyShaderPath, "miss");
            /*
            if (lightPrefix == LightSampling.envmap)
            {
                Context["use_envmap"].Set(1);
                Program envMapMiss = new Program(Context, shaderPath, "envmap_miss");
                var envmap = ImageFactory.CreateTextureSampler(Context,
                                                               @"F:\3d\HDRi\Harbor_3_Free_Ref.hdr",
                                                               true); //turistskaya_pan.jpg");//KC_outside_hi.hdr
                Context["envmap"].Set(envmap);
                Context["envmap_lt"].Set(ImageFactory.CreateTextureSampler(Context,
                                                                           @"F:\3d\HDRi\Harbor_3_Free_Env.hdr",
                                                                           true));
                Context.SetRayMissProgram(0, envMapMiss);

            }
            else
                if (lightPrefix == LightSampling.sun || !hasLights)
                {
                    Context.SetRayMissProgram(0, skyMiss);
                }
                else
                {
                    Context.SetRayMissProgram(0, miss);
                }


            */
            //END SETUP SKY

            if (Scene.EnvMap != null)
            {
                Context["use_envmap"].Set(1);
                Program envMapMiss = new Program(Context, shaderPath, "envmap_miss");
                var envmap = Scene.EnvMap.EnvironmentMap;
                Context["envmap"].Set(envmap);
                Context["envmap_lt"].Set(envmap);
                Context.SetRayMissProgram(0, envMapMiss);
            }
            else
            {
                Context.SetRayMissProgram(0, miss);

                var dm = EngineContext.Instance.TexManager.CreateTextureSampler(@"F:\3D\HDRI\empty.jpg");
                Context["envmap"].Set(dm);
                Context["envmap_lt"].Set(dm);

            }

            Context.SetRayGenerationProgram(0, rayGen);


            miss["bg_color"].Set(10 / 255.0f, 14 / 255.0f, 23 / 255.0f);

            Context.SetExceptionProgram(0, exception);

            /*-----------------------------------------------
             * Finally compile the optix context, and build the accel tree
             *-----------------------------------------------*/
            setCamera(Scene.BBox);
            //skyModel._up = this.Camera.Up;

            //skyModel.InitContext(Context);
            Context["directional_light"].SetFloat3(skyModel._sun_dir * 100000f);
            Context["directional_light_col"].SetFloat3(skyModel._sun_color * (1.0f / 1600000.0f));
            Context["directional_light_col"].SetFloat3(skyModel._sun_color * (1.0f / 1600000.0f));
            //Context["emission_color"].SetFloat3(skyModel._sun_color * (1.0f / 1600000.0f));
            //Context["emission_color"].SetFloat3(new Vector3(0,0,0));


            Context["top_object"].Set(this.sceneManager.Scene.TopObject);
            Context["output_buffer"].Set(OutputBuffer);
            Context["scene_epsilon"].Set(this.SceneEpsilon);


            Context["rr_begin_depth"].Set(rrBeginDepth);
            Context["max_depth"].Set(maxDepth);

            Context["sqrt_num_samples"].Set(sqrtSamples);
            Context["frame_number"].Set(mFrame);

            Context["pathtrace_ray_type"].Set(0u);
            Context["pathtrace_shadow_ray_type"].Set(1u);
            Context["glass_color"].SetFloat3(new Vector3(0.50f, 0.50f, 0.50f));
            Context["diffuse_color"].SetFloat3(new Vector3(0.50f, 0.50f, 0.50f));
            Context["shadow_attenuation"].SetFloat3(new Vector3());
            Context["extinction_constant"].SetFloat3(new Vector3() + 0.75f);


            Context["index_of_refraction"].Set(1.7f);
            Context["focal_scale"].Set(0.0f);
            Context["aperture_radius"].Set(apertureSize);



            Context["bad_color"].SetFloat3(new Vector3(1.0f, 0.0f, 1.0f));

            Console.Write("Compiling Optix... ");
            float start = Time.GetTimeInSecs();
            //if (!LoadAccelData(model, modelPath, accelCache)) {
            //    model.GeoGroup.Acceleration.MarkAsDirty();
            //}

            // bool accelIsDirty = model.GeoGroup.Acceleration.IsDirty();

            Context.Compile();

            Context.BuildAccelTree();

            Console.WriteLine("{0:0.00}s", Time.GetTimeInSecs() - start);
        }

        private void CreateLights(IEnumerable<TriangleLight> lights)
        {

            BufferDesc desc = new BufferDesc() { Width = (uint)(lights.Count()), Format = Format.User, Type = BufferType.Input, ElemSize = (uint)Marshal.SizeOf(typeof(TriangleLight)) };
            OptixDotNet.Buffer lightBuffer = new OptixDotNet.Buffer(Context, desc);
            BufferStream stream = lightBuffer.Map();
            {
                foreach (var triangleLight in lights)
                {
                    triangleLight.normal.Normalize();
                    stream.Write(triangleLight);
                }
            }
            lightBuffer.Unmap();

            Context["lights"].Set(lightBuffer);
        }

        private Dictionary<string, Material> InitializeMaterials(Context ctx, string shaderPath, string sh2 = null)
        {
            Console.WriteLine("Materials");
            var matLib = new Dictionary<string, Material>();

            var diffuseMat = new Material(ctx);
            diffuseMat.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "diffuse_" + lightPrefix);
            diffuseMat.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");
            matLib.Add("diffuse", diffuseMat);

            var diffuseTexMat = new Material(ctx);
            diffuseTexMat.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "diffuse_tex_" + lightPrefix);
            diffuseTexMat.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");
            matLib.Add("diffuse_tex", diffuseTexMat);


            var specularMat = new Material(ctx);
            specularMat.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, sh2, "specular_glass");
            specularMat.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "glass_shadow");
            matLib.Add("specular", specularMat);


            var diffuseEmitter = new Material(ctx);
            diffuseEmitter.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "diffuseEmitterCL");
            diffuseEmitter.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "diffuseEmitterShadow");
            matLib.Add("diffuseEmitter", diffuseEmitter);


            var metal = new Material(ctx);
            metal.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "metal_" + lightPrefix);
            metal.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");
            matLib.Add("metal", metal);


            var glossy_metal = new Material(ctx);
            glossy_metal.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "glossymetal_" + lightPrefix);
            glossy_metal.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");
            matLib.Add("glossy_metal", glossy_metal);

            var metalTex = new Material(ctx);
            metalTex.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "metal_tex_" + lightPrefix);
            metalTex.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");
            matLib.Add("metal_tex", metalTex);


            var alloy = new Material(ctx);
            alloy.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "diffuse_" + lightPrefix);
            alloy.Programs[1] = new SurfaceProgram(Context, RayHitType.Any, shaderPath, "shadow");
            matLib.Add("alloy", alloy);

            //Setup Glass

            //var glass = new Material(ctx);

            //glass.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "glass");
            //glass.Programs[0] = new SurfaceProgram(Context, RayHitType.Closest, shaderPath, "glass_shadow");

            return matLib;
        }



        public override void UpdateFrame()
        {
            base.UpdateFrame();
            Context["Y_log_av"].Set(yAvg);
            Context["Y_max"].Set(yMax);
            skyModel.InitContext(Context);
        }
    }
}
